#!/usr/bin/env python3
"""
IMU → 世界系点云 → 帧分批 matplotlib 3D 动画
python3.12 bag2world_pcl.py \
  --bag unilidar-2023-09-22-12-42-04.db3 \
  --bunch 200
"""
import argparse, numpy as np, matplotlib.pyplot as plt
from matplotlib import animation
from rosbag2_py import SequentialReader, StorageOptions, ConverterOptions, StorageFilter
from sensor_msgs.msg import PointCloud2, Imu
from rclpy.serialization import deserialize_message
from functools import partial

# ---------- 工具 ----------
def quat2R(qw, qx, qy, qz):
    """四元数 → 3×3 旋转矩阵"""
    R = np.array([
        [1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
        [2 * (qx * qy + qw * qz), 1 - 2 * (qx ** 2 + qz ** 2), 2 * (qy * qz - qw * qx)],
        [2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx ** 2 + qy ** 2)]
    ], dtype=np.float32)
    return R

def pc2_to_xyz(pc: PointCloud2):
    """解析 PointCloud2 → (N,3) np.float32，过滤无效点"""
    arr = np.frombuffer(pc.data, np.uint8).reshape(-1, pc.point_step)
    xyz = arr[:, :12].copy().view(np.float32).reshape(-1, 3)
    return xyz[np.isfinite(xyz).all(axis=1)]

# ---------- 主逻辑 ----------
def load_world_bunches(bag_path, bunch_size):
    reader = SequentialReader()
    reader.open(StorageOptions(uri=bag_path, storage_id='sqlite3'),
                ConverterOptions('', ''))
    # 同时过滤两个 topic，保证时间顺序
    reader.set_filter(StorageFilter(topics=['/unilidar/imu', '/unilidar/cloud']))

    bunches, cache = [], []
    latest_R = np.eye(3, dtype=np.float32)   # 初始姿态
    while reader.has_next():
        topic, data, t = reader.read_next()
        if topic == '/unilidar/imu':
            imu = deserialize_message(data, Imu)
            q = imu.orientation
            latest_R = quat2R(q.w, q.x, q.y, q.z)
        elif topic == '/unilidar/cloud':
            pc = deserialize_message(data, PointCloud2)
            pts_imu = pc2_to_xyz(pc)
            pts_world = (latest_R @ pts_imu.T).T          # 3×3 · 3×N  → 世界系
            cache.append(pts_world)

            if len(cache) == bunch_size:                  # 满一批
                bunches.append(np.vstack(cache))          # 仅这批
                cache = []
    if cache:                                             # 尾部
        bunches.append(np.vstack(cache))
    return bunches

# ---------- 动画 ----------
def animate(i, bunches):
    pts = bunches[i]
    print(f'step {i+1}/{len(bunches)}  points={len(pts)}')
    scat._offsets3d = (pts[:, 0], pts[:, 1], pts[:, 2])
    ax.set_title(f'step {i+1}/{len(bunches)}  pts={len(pts)}')
    # 动态坐标轴
    margin = 1.0
    mins, maxs = pts.min(axis=0), pts.max(axis=0)
    ax.set_xlim(mins[0] - margin, maxs[0] + margin)
    ax.set_ylim(mins[1] - margin, maxs[1] + margin)
    ax.set_zlim(mins[2] - margin, maxs[2] + margin)
    return scat,

# ---------- main ----------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bag', default='/gitlab/pcd_analysis/demo_data/unilidar-2023-09-22-12-42-04.db3', help='ros2 bag db3 path')
    parser.add_argument('--bunch', type=int, default=20, help='frames per step')
    parser.add_argument('--interval', type=int, default=50, help='ms per step')
    args = parser.parse_args()

    print('读取 bag 并按 IMU 转到世界系 …')
    bunches = load_world_bunches(args.bag, args.bunch)
    print(f'共 {len(bunches)} 步，准备播放')

    fig = plt.figure(figsize=(8, 6))
    global ax, scat
    ax = fig.add_subplot(111, projection='3d')
    scat = ax.scatter([], [], [], s=1)
    ax.set_box_aspect([1, 1, 1])
    ax.set_xlabel('X world')
    ax.set_ylabel('Y world')
    ax.set_zlabel('Z world')

    ani = animation.FuncAnimation(fig, partial(animate, bunches=bunches), frames=len(bunches),
                                  interval=args.interval, blit=False, repeat=True)
    plt.show()

if __name__ == '__main__':
    main()